package com.google.p2pbench;

import android.content.BroadcastReceiver;
import android.content.Context;
import android.content.Intent;
import android.content.IntentFilter;
import android.net.NetworkInfo;
import android.net.wifi.WpsInfo;
import android.net.wifi.p2p.WifiP2pConfig;
import android.net.wifi.p2p.WifiP2pDevice;
import android.net.wifi.p2p.WifiP2pInfo;
import android.net.wifi.p2p.WifiP2pManager;
import android.net.wifi.p2p.WifiP2pManager.ActionListener;
import android.net.wifi.p2p.WifiP2pManager.Channel;
import android.net.wifi.p2p.WifiP2pManager.ChannelListener;
import android.net.wifi.p2p.WifiP2pManager.ConnectionInfoListener;
import android.net.wifi.p2p.WifiP2pManager.DnsSdServiceResponseListener;
import android.net.wifi.p2p.WifiP2pManager.DnsSdTxtRecordListener;
import android.net.wifi.p2p.nsd.WifiP2pDnsSdServiceInfo;
import android.net.wifi.p2p.nsd.WifiP2pDnsSdServiceRequest;
import android.os.Handler;
import android.os.Looper;
import android.os.Parcel;
import android.os.Parcelable;
import android.util.Log;
import android.util.Pair;

import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;
import java.net.DatagramPacket;
import java.net.DatagramSocket;
import java.net.InetAddress;
import java.net.InetSocketAddress;
import java.net.SocketAddress;
import java.nio.ByteBuffer;
import java.util.HashMap;
import java.util.HashSet;
import java.util.Map;
import java.util.Random;
import java.util.Set;

public class WifiDirectConnectivity implements Connectivity, ChannelListener {
  private static final String TAG = "WifiDirectConnectivity";

  private final Context mContext;
  private final WifiP2pManager mWifiP2pManager;
  private Channel mChannel;
  private WifiStateObserver mWifiStateObserver;

  private DatagramSocket mDatagramSocket;
  private WifiP2pInfo mWifiInfo;

  private Listener mListener;
  private final Map<Pair<String, Service>, ServiceBackend> mServices =
      new HashMap<Pair<String, Service>, ServiceBackend>();
  private final Map<Pair<String, Discoverer>, DiscovererBackend> mDiscoverers =
      new HashMap<Pair<String, Discoverer>, DiscovererBackend>();

  private final Handler mHandler = new Handler();

  private class WifiStateObserver extends BroadcastReceiver {
    @Override
    public void onReceive(Context context, final Intent intent) {
      String action = intent.getAction();
      Log.w(TAG, "received " + action);
      if (!WifiP2pManager.WIFI_P2P_CONNECTION_CHANGED_ACTION.equals(action))
        return;

      NetworkInfo networkInfo = (NetworkInfo) intent
          .getParcelableExtra(WifiP2pManager.EXTRA_NETWORK_INFO);

      if (!networkInfo.isConnected()) {
        mHandler.post(new Runnable() {
          @Override
          public void run() {
            onWifiChange(null);
          }
        });
        return;
      }

      mWifiP2pManager.requestConnectionInfo(mChannel,
        new ConnectionInfoListener() {
          @Override
          public void onConnectionInfoAvailable(final WifiP2pInfo info) {
            onWifiChange(info);
          }
        }
      );

    }
  }

  private void registerObserver() {
    if (mWifiStateObserver != null) return;
    mWifiStateObserver = new WifiStateObserver();
    IntentFilter filter = new IntentFilter();
    // filter.addAction(WifiP2pManager.WIFI_P2P_STATE_CHANGED_ACTION);
    filter.addAction(WifiP2pManager.WIFI_P2P_CONNECTION_CHANGED_ACTION);
    mContext.registerReceiver(mWifiStateObserver, filter);
  }

  private void unregisterObserver() {
    if (mWifiStateObserver == null) return;
    mContext.unregisterReceiver(mWifiStateObserver);
    mWifiStateObserver = null;
  }

  private void onWifiChange(WifiP2pInfo info) {
    Log.w(TAG, "onWifiChange " + info);
    synchronized (this) {
      mWifiInfo = info;
      notify();
    }
    // TODO: if info == null or info.groupFormed = false, then we should restart discovery?

    // OK, this is totally bizarre, but...
    // Once a group is formed, discovery stops, so make sure to resume discovery
    // after group is disbanded.
    if (info == null) {
      mP2pDiscoverer.start();
    } else {
      mP2pDiscoverer.stop();
    }
  }

  // There can be only one DnsSdServiceResponseListener, so this is going to be it.
  private class P2pDiscoverer implements DnsSdServiceResponseListener, DnsSdTxtRecordListener {
    private final Map<String, Set<DiscovererBackend>> mServiceTypes =
        new HashMap<String, Set<DiscovererBackend>>();
    private boolean mDiscovering = false;

    WifiP2pManager.ActionListener mActionListener = new WifiP2pManager.ActionListener() {
      @Override
      public void onSuccess() {
        Log.d(TAG, "Success");
      }

      @Override
      public void onFailure(int errorCode) {
        // TODO: communicate errors somehow
        // mDiscoverer.onError(new ConnectivityException("Discovery failed " + errorCode));
        Log.e(TAG, "Error " + errorCode);
      }
    };

    @Override
    public void onDnsSdServiceAvailable(
        String instanceName, String serviceType, WifiP2pDevice device) {
      // TODO: save as discovered services??
      Log.w(TAG, "onDnsSdServiceAvailable " + instanceName + " " + serviceType);
    }

    @Override
    public void onDnsSdTxtRecordAvailable(
        String fullDomainName, Map<String, String> record, WifiP2pDevice device) {
      Log.w(TAG, "onDnsSdTxtRecordAvailable " + fullDomainName + " " + record);
      // TODO: save port for each service?
      int dot = fullDomainName.indexOf('.');
      if (dot < 0) return;
      String serviceType = fullDomainName.substring(dot + 1);
      // Strip .local. suffix automatically appended by WifiP2pDnsSdServiceInfo.
      final String SUFFIX = ".local.";
      if (serviceType.endsWith(SUFFIX))
        serviceType = serviceType.substring(0, serviceType.length() - SUFFIX.length());
      try {
        int port = Integer.parseInt(record.get("port"));
        notifyDiscoverers(serviceType, device, port);
      } catch (java.lang.NumberFormatException e) {
        // ignore this service
      }
    }

    private void notifyDiscoverers(String serviceType, WifiP2pDevice device, int port) {
      Log.w(TAG, "notifyDiscoverers " + serviceType);
      Set<DiscovererBackend> backends = mServiceTypes.get(serviceType);
      if (backends != null) {
        for (DiscovererBackend b : backends)
          b.onDiscovery(device, port);
      }
    }

    void register(String serviceType, DiscovererBackend backend) {
      Log.w(TAG, "discovering " + serviceType);
      // NOTE: WifiP2pManager automatically appends .local. to all service requests
      // frameworks/base/wifi/java/android/net/wifi/p2p/nsd/WifiP2pDnsSdServiceInfo.java#125
      // but does not strip it.
      Set<DiscovererBackend> backends = mServiceTypes.get(serviceType);
      if (backends == null) {
        backends = new HashSet<DiscovererBackend>(1); // Minimal size.
        mServiceTypes.put(serviceType, backends);
        // NOTE: If we specify |serviceType| but not |serviceName|, we won't receive
        // TXT records at all. Instead let's ask for all services and do it in start().
        // mWifiP2pManager.addServiceRequest(
        // mChannel, WifiP2pDnsSdServiceRequest.newInstance(serviceType), mActionListener);
      }
      backends.add(backend);

      // TODO: try to serve from existing discoveries????
      start();
    }

    void unregister(String serviceType, DiscovererBackend backend) {
      Set<DiscovererBackend> backends = mServiceTypes.get(serviceType);
      if (backends == null) return;
      backends.remove(backend);
      if (backends.isEmpty()) {
        mServiceTypes.remove(serviceType);
        Log.w(TAG, "removeServiceRequest " + serviceType);
      }
      if (mServiceTypes.isEmpty()) stop();
    }

    void start() {
      if (mDiscovering) return;
      mDiscovering = true;
      // NOTE: if we have no service discovery requests, then we won't be discovering.
      // See register() for why we ask for all services anyway.
      mWifiP2pManager.addServiceRequest(
          mChannel, WifiP2pDnsSdServiceRequest.newInstance(), mActionListener);
      // TODO: handle action result
      Log.w(TAG, "setDnsSdResponseListeners");
      mWifiP2pManager.setDnsSdResponseListeners(mChannel, this, this);
      Log.w(TAG, "discoverServices");
      mWifiP2pManager.discoverServices(mChannel, mActionListener);
    }

    void stop() {
      if (!mDiscovering) return;
      mDiscovering = false;
      // TODO: handle action result
      Log.w(TAG, "clearServiceRequests");
      mWifiP2pManager.clearServiceRequests(mChannel, mActionListener);
    }
  }

  private final P2pDiscoverer mP2pDiscoverer = new P2pDiscoverer();

  private static class PeerImpl implements Peer {
    final WifiP2pDevice mDevice;
    final SocketAddress mAddress;

    PeerImpl(WifiP2pDevice device, InetAddress address, int port) {
      mDevice = device;
      mAddress = new java.net.InetSocketAddress(address, port);
    }

    PeerImpl(WifiP2pDevice device, SocketAddress address) {
      mDevice = device;
      mAddress = address;
    }

    @Override
    public String getDescription() {
      return "[WifiDirect " + mDevice + " " + mAddress + "]";
    }

    @Override
    public int describeContents() {
      return 0;
    }

    @SuppressWarnings("unused")
    public static final Parcelable.Creator<PeerImpl> CREATOR = new Parcelable.Creator<PeerImpl>() {
      @Override
      public PeerImpl createFromParcel(Parcel in) {
        return new PeerImpl(
            WifiP2pDevice.CREATOR.createFromParcel(in), (SocketAddress) in.readSerializable());
      }

      @Override
      public PeerImpl[] newArray(int size) {
        return new PeerImpl[size];
      }
    };

    @Override
    public void writeToParcel(Parcel dest, int flags) {
      dest.writeParcelable(mDevice, flags);
      dest.writeSerializable(mAddress);
    }
  }

  private static class SocketImpl implements Socket {
    java.net.Socket mSocket;

    SocketImpl(java.net.Socket socket) {
      mSocket = socket;
    }

    @Override
    public void close() {
      try {
        mSocket.close();
      } catch (java.io.IOException e) {
        // ignore
      }
    }

    @Override
    public boolean isConnected() {
      return mSocket.isConnected();
    }

    @Override
    public InputStream getInputStream() throws ConnectivityException {
      try {
        return mSocket.getInputStream();
      } catch (java.io.IOException e) {
        throw new ConnectivityException("Socket failed", e);
      }
    }

    @Override
    public OutputStream getOutputStream() throws ConnectivityException {
      try {
        return mSocket.getOutputStream();
      } catch (java.io.IOException e) {
        throw new ConnectivityException("Socket failed", e);
      }
    }
  }

  private class ServiceBackend {
    final String mServiceType;
    final Service mService;

    java.net.ServerSocket mServerSocket;
    java.net.DatagramSocket mMessageSocket;
    Thread mAcceptorThread;
    Thread mDatagramThread;

    WifiP2pDnsSdServiceInfo mServiceInfo;

    WifiP2pManager.ActionListener mRegistrationListener = new WifiP2pManager.ActionListener() {
      @Override
      public void onSuccess() {
        Log.d(TAG, "Success");
      }
      @Override
      public void onFailure(int errorCode) {
        Log.e(TAG, "Error " + errorCode);
        mService.onError(new ConnectivityException("Registration failed " + errorCode));
      }
    };

    Runnable mAcceptor = new Runnable() {
      @Override
      public void run() {
        try {
          while (true) {
            java.net.Socket socket = mServerSocket.accept();
            mService.onStream(
                new PeerImpl(null, socket.getRemoteSocketAddress()),
                new SocketImpl(socket));
          }
        } catch (IOException e) {
          mService.onError(new ConnectivityException("Acceptor failed", e));
          shutdown();
        }
      }
    };

    Runnable mMessageReader = new Runnable() {
      @Override
      public void run() {
        try {
          byte[] data = new byte[2500];
          DatagramPacket packet = new DatagramPacket(data, data.length);
          while (true) {
            mMessageSocket.receive(packet);
            ByteBuffer msg = ByteBuffer.wrap(packet.getData(), packet.getOffset(),
                packet.getLength());
            mService.onMessage(new PeerImpl(null, packet.getSocketAddress()), msg);
          }
        } catch (IOException e) {
          mService.onError(new ConnectivityException("DatagramReader failed", e));
        }
      }
    };

    ServiceBackend(String serviceType, Service service) throws ConnectivityException {
      mServiceType = serviceType;
      mService = service;

      for (int i = 0; i < 4; ++i) {
        try {
          // One port for both UDP and TCP so that there's only one DNS-SD announcement.
          int port = (new Random()).nextInt(1024) + (65536 - 1024);
          mMessageSocket = new DatagramSocket(port);
          mServerSocket = new java.net.ServerSocket(port);
        } catch (IOException e) {
          // try again
          if (mServerSocket != null) {
            try {
              mServerSocket.close();
            } catch (IOException f) {
              // ignore
            }
            mServerSocket = null;
          }
          if (mMessageSocket != null) {
            mMessageSocket.close();
            mMessageSocket = null;
          }
        }
      }

      if (mServerSocket == null) {
        throw new ConnectivityException("Failed to start service");
      }

      mAcceptorThread = new Thread(mAcceptor, mService.toString() + " acceptor");
      mAcceptorThread.start();
      mDatagramThread = new Thread(mMessageReader, mService.toString() + " datagram reader");
      mDatagramThread.start();

      Map<String, String> records = new HashMap<String, String>();
      // NOTE: we don't know our InetAddress yet!
      records.put("port", String.valueOf(mServerSocket.getLocalPort()));
      String uid = "_" + Long.toHexString((new Random()).nextLong());
      mServiceInfo = WifiP2pDnsSdServiceInfo.newInstance(uid, mServiceType, records);
      Log.w(TAG, "addLocalService " + mServiceInfo);
      mWifiP2pManager.addLocalService(mChannel, mServiceInfo, mRegistrationListener);

      // Seriously? https://code.google.com/p/android/issues/detail?id=37425
      mP2pDiscoverer.start();
    }

    PeerImpl getPeer() {
      // TODO: get my WifiP2pDevice
      // TODO: only port of the local SocketAddress is meaningful.
      return new PeerImpl(null, mServerSocket.getLocalSocketAddress());
    }

    synchronized void shutdown() {
      if (mServerSocket != null) {
        mWifiP2pManager.removeLocalService(mChannel, mServiceInfo, mRegistrationListener);
        try {
          mServerSocket.close();
        } catch (IOException e) {
          // ignore
        }
        mServerSocket = null;
      }
      if (mMessageSocket != null) {
        mMessageSocket.close();
        mMessageSocket = null;
      }
    }
  }

  private class DiscovererBackend {
    String mServiceType;
    Discoverer mDiscoverer;

    DiscovererBackend(String serviceType, Discoverer discoverer) {
      mServiceType = serviceType;
      mDiscoverer = discoverer;

      mP2pDiscoverer.register(mServiceType, this);
    }

    void onDiscovery(WifiP2pDevice srcDevice, int port) {
      Peer peer = new PeerImpl(srcDevice, null, port);
      mDiscoverer.onDiscovery(mServiceType, peer);
    }

    void shutdown() {
      mP2pDiscoverer.unregister(mServiceType, this);
    }
  }


  public WifiDirectConnectivity(Context context) {
    mContext = context;
    mWifiP2pManager = (WifiP2pManager) context.getSystemService(Context.WIFI_P2P_SERVICE);
  }

  @Override
  public boolean supportsMessages() {
    return true;
  }

  @Override
  public boolean connect(Listener listener) {
    Util.assertOnThread(mHandler);
    mListener = listener;

    // ((WifiManager) mContext.getSystemService(Context.WIFI_SERVICE)).setWifiEnabled(true);
    registerObserver();

    mChannel = mWifiP2pManager.initialize(mContext, Looper.getMainLooper(), this);
    if (mListener != null) mListener.onConnected();
    return true;
  }

  @Override
  public void shutdown(int mode) {
    Util.assertOnThread(mHandler);
    if (mDatagramSocket != null)
      mDatagramSocket.close(); // this kills the Receiver thread
    mDatagramSocket = null;
    unregisterObserver();
    mWifiP2pManager.removeGroup(mChannel, null);
    for (ServiceBackend b : mServices.values()) {
      b.shutdown();
    }
    for (DiscovererBackend b : mDiscoverers.values()) {
      b.shutdown();
    }
    mP2pDiscoverer.stop();
    // mWifiManager.setWifiEnabled(false);
    if (mListener != null)
      mListener.onError(new ConnectivityException("shutdown"));
  }

  @Override
  public Peer getBroadcastPeer() {
    return null;
  }

  @Override
  public void setBroadcastService(Service service) throws ConnectivityException {
    throw new ConnectivityException("Not supported");
  }

  @Override
  public void sendMessage(Peer peer, ByteBuffer msg) throws ConnectivityException {
    throw new ConnectivityException("Not yet implemented");
  }

  @Override
  public Socket getSocket(Peer peer) throws ConnectivityException {
    synchronized (this) {
      while (mWifiInfo == null) {
        PeerImpl peerImpl = (PeerImpl) peer;
        WifiP2pConfig config = new WifiP2pConfig();
        config.deviceAddress = peerImpl.mDevice.deviceAddress;
        config.groupOwnerIntent = 0; // least inclination to be a group owner
        config.wps.setup = WpsInfo.PBC; // this is default
        Log.w(TAG, "getSocket connect " + config.deviceAddress);
        mWifiP2pManager.connect(mChannel, config, new ActionListener() {
          @Override
          public void onSuccess() {
            Log.d(TAG, "connect success");
          }

          @Override
          public void onFailure(int code) {
            Log.e(TAG, "connect fail " + code);
            // TODO: must throw connectivityexception from getSocket
          }
        });
        try {
          wait();
        } catch (InterruptedException e) {
          Log.e(TAG, "interrupted while waiting for connection info");
          return null;
        } // until connected
      }
    }
    int port = ((InetSocketAddress) ((PeerImpl) peer).mAddress).getPort();
    java.net.Socket socket = new java.net.Socket();
    try {
      socket.connect(new InetSocketAddress(mWifiInfo.groupOwnerAddress, port));
    } catch (IOException e) {
      try {
        socket.close();
      } catch (IOException f) {
        // ignore
      }
      throw new ConnectivityException("Could not connect", e);
    }
    return new SocketImpl(socket);
  }

  @Override
  public Peer announce(String serviceType, Service service) throws ConnectivityException {
    Util.assertOnThread(mHandler);
    ServiceBackend backend = new ServiceBackend(serviceType, service);
    mServices.put(Pair.create(serviceType, service), backend);
    return backend.getPeer();
  }

  @Override
  public void unnannounce(String serviceType, Service service) {
    Util.assertOnThread(mHandler);
    ServiceBackend backend = mServices.remove(Pair.create(serviceType, service));
    if (backend != null) backend.shutdown();
  }

  @Override
  public void discover(String serviceType, Discoverer listener) {
    Util.assertOnThread(mHandler);
    DiscovererBackend backend = new DiscovererBackend(serviceType, listener);
    mDiscoverers.put(Pair.create(serviceType, listener), backend);
  }

  @Override
  public void stopDiscovery(String serviceType, Discoverer listener) {
    Util.assertOnThread(mHandler);
    DiscovererBackend backend = mDiscoverers.remove(Pair.create(serviceType, listener));
    if (backend != null) backend.shutdown();
  }

  @Override
  public void onChannelDisconnected() {
    if (mListener != null) mListener.onError(new ConnectivityException("channel disconnected"));
  }
}
